/*
 * Copyright 2021 The happyProxy Project
 *
 * The happyProxy Project licenses this file to you under the Apache License,
 * version 2.0 (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at:
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */
package com.sinosoft.lingkang.happyproxy.server;

import com.sinosoft.lingkang.happyproxy.HappyProxyUtils;
import com.sinosoft.lingkang.happyproxy.constants.HappyProxyConstants;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelOption;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpRequest;
import io.netty.util.CharsetUtil;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;

/**
 * @author 绫小路
 * @date 2021/2/12 0:50
 * @description
 */
public class ConnectToRemoteServerHandler extends ChannelInboundHandlerAdapter {

  private static final InternalLogger logger = InternalLoggerFactory.getInstance(ConnectToRemoteServerHandler.class);
  // As we use inboundChannel.eventLoop() when building the Bootstrap this does not need to be volatile as
  // the outboundChannel will use the same EventLoop (and therefore Thread) as the inboundChannel.
  private Channel outboundChannel;

  private HttpAuthorization authorization;
  private HttpProxyServerRequestFilter requestFilter;
  private HttpProxyServerFlowCount httpProxyServerFlowCount;


  public ConnectToRemoteServerHandler setHttpProxyServerFlowCount(HttpProxyServerFlowCount httpProxyServerFlowCount) {
    this.httpProxyServerFlowCount = httpProxyServerFlowCount;
    return this;
  }

  protected ConnectToRemoteServerHandler setAuthorization(HttpAuthorization authorization) {
    this.authorization = authorization;
    return this;
  }

  protected ConnectToRemoteServerHandler setRequestFilter(HttpProxyServerRequestFilter requestFilter) {
    this.requestFilter = requestFilter;
    return this;
  }

  //----------------------------------------------------------------------------------------------------------------

  @Override
  public void channelRead(final ChannelHandlerContext ctx, Object msg) throws Exception {
    if (outboundChannel != null && outboundChannel.isActive()) {
      outboundChannel.writeAndFlush(msg).addListener(new ChannelFutureListener() {
        @Override
        public void operationComplete(ChannelFuture future) {
          if (future.isSuccess()) {
            // was able to flush out data, start to read the next chunk
            ctx.channel().read();
          } else {
            future.channel().close();
          }
        }
      });
    } else if (msg instanceof HttpRequest) {
      final Channel inboundChannel = ctx.channel();
      final HttpRequest request = (HttpRequest) msg;
      String[] hosts = request.headers().get("Host").split(":");
      String host = null;
      int port = 80;
      if (hosts.length == 1) {
        host = hosts[0].replace("\r", "");
      } else {
        host = hosts[0];
        port = Integer.parseInt(hosts[1].replace("\r", ""));
      }
      //认证
      if (authorization != null) {
        String proxyAuth = request.headers().get(HappyProxyConstants.HTTP_PROXY_AUTHORIZATION_HEADER);
        if (HappyProxyUtils.isEmpty(proxyAuth)) {
          //向客户端返回一个407，告诉客户端需要账号密码
          inboundChannel.writeAndFlush(Unpooled.wrappedBuffer(HappyProxyConstants.PROXY_AUTHORIZATION_FAIL_RESPONSE));
          super.channelRead(ctx, msg);
          return;
        }
        String[] account = HappyProxyUtils.proxyAuthDecoder(proxyAuth.substring(6)).split(":");
        if (account.length == 1) {
          //向客户端返回一个407，告诉客户端需要账号密码
          inboundChannel.writeAndFlush(Unpooled.wrappedBuffer(HappyProxyConstants.PROXY_AUTHORIZATION_FAIL_RESPONSE));
          super.channelRead(ctx, msg);
          return;
        }
        if (!authorization.authorization(account[0], account[1])) {
          //向客户端返回一个407，告诉客户端需要账号密码
          inboundChannel.writeAndFlush(Unpooled.wrappedBuffer(HappyProxyConstants.PROXY_AUTHORIZATION_FAIL_RESPONSE));
          super.channelRead(ctx, msg);
          return;
        } else {
          //proxy security
          request.headers().remove(HappyProxyConstants.HTTP_PROXY_AUTHORIZATION_HEADER);
        }
      }
      //过滤
      if (requestFilter != null) {
        if (!requestFilter.doFilter(request, inboundChannel, host, port)) {
          super.channelRead(ctx, msg);
          return;
        }
      }

      // Start the connection attempt.
      Bootstrap b = new Bootstrap();
      b.group(inboundChannel.eventLoop())
          .channel(ctx.channel().getClass())
          .handler(new RemoteServerResponseHandler(inboundChannel, httpProxyServerFlowCount))
          .option(ChannelOption.AUTO_READ, false);
      ChannelFuture f = b.connect(host, port);
      outboundChannel = f.channel();
      f.addListener(new ChannelFutureListener() {
        @Override
        public void operationComplete(ChannelFuture future) {
          if (future.isSuccess()) {
            if (request.method().equals(HttpMethod.CONNECT)) {
              //Https 升级为https需要与client握手确认一次
              inboundChannel.writeAndFlush(Unpooled.copiedBuffer("HTTP/1.1 200 Connection Established\r\n\r\n", CharsetUtil.UTF_8));
            } else {
              //如果是http连接，直接将数据发送到远端服务
              future.channel().writeAndFlush(HappyProxyUtils.httpRequestToByteBuf(request));
            }
            ctx.pipeline().remove(HappyProxyConstants.HTTP_REQUEST_DECODER);

            // connection complete start to read first data
            inboundChannel.read();
          } else {
            // Close the connection if the connection attempt has failed.
            inboundChannel.close();
          }
        }
      });

    }
  }

  @Override
  public void channelInactive(ChannelHandlerContext ctx) {
    if (outboundChannel != null) {
      closeOnFlush(outboundChannel);
    }
  }

  @Override
  public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
    logger.error(cause);
    closeOnFlush(ctx.channel());
  }

  /**
   * Closes the specified channel after all queued write requests are flushed.
   */
  private void closeOnFlush(Channel ch) {
    if (ch.isActive()) {
      ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
    }
  }

  //-----------------------------------------------------------------------------------------

}
